import os
import random
import sys
from enum import Enum
import psutil
import shutil

CLUS_SIZE = 0x20
ROOT_START_IDX = 16  # 以M为单位
RESERVE_SIZE = 0x5000  # 保留扇区数 10M/512

class SPEED_FILL(Enum):
    FAST = 1 # 只将第4个簇填充为未使用
    MID = 2 # 只将前100M或小于100M的FAT填充为未使用
    SLOW = 3 # 将整个FAT填充为未使用


def fill_fat(fd, fat_size: int,tp):
    
    data = int(0x0ffffff8).to_bytes(4, 'little')
    data = data + int(0xffffffff).to_bytes(4, 'little')
    data = data + int(0x0fffffff).to_bytes(4, 'little')
    if tp==SPEED_FILL.FAST:
        data = data+int(0).to_bytes(4,'little')
        os.write(fd, data)
        return
    data = data + (b'\x00' * 500)
    os.write(fd,data)
    fat_size -= 1
    i_count=0

    while True:
        if fat_size <= 0:
            break
        
        tmp_size = fat_size - 2048*10
        if (tmp_size < 0):
            tmp_size = fat_size
        else:
            tmp_size=2048*10
        data = b'\x00' * 512 * tmp_size
        os.write(fd, data)
        fat_size -= tmp_size
        i_count+=1
        # if (i_count%10==0):
        print('已初始化{}M'.format(i_count*10))

def fill_root_dir(fd):
    data = b'HWCH    ' # 文件名
    data = data + b'   ' # 扩展名
    data = data + b'\x08' #属性字节
    data = data + b'\x00' #系统保留
    data = data + b'\x00' # 创建时间的10毫秒
    data = data + b'\x00\x00' # 文件创建时间
    data = data + b'\x00\x00' #文件创建日期
    data = data + b'\x00\x00' # 文件最后访问时间
    data = data + b'\x00\x00' # 文件起始簇号高16位
    data = data + b'\x00\x00' # 文件最近修改时间
    data = data + b'\x00\x00' # 文件最近修改日期
    data = data + b'\x00\x00' # 文件起始簇号低16位
    data = data + b'\x00\x00\x00\x00' #文件长度
    os.write(fd, data)
    size = CLUS_SIZE * 512 - 32
    data = b'\x00' * size
    os.write(fd, data)

def is_mount(diskpath):
    for part in psutil.disk_partitions():
        if part.device.find(diskpath) != -1:
            return True

    return False

def write_dos(dos_path: str, tp):
    is_dev=False
    if (dos_path.startswith('/dev/')):
        if sys.argv[1][-1].isdigit():
            dos_path=sys.argv[1][:-1]
        else:
            dos_path=sys.argv[1]
        if dos_path[-1].isdigit():
            dos_path=dos_path[:-1]

        with open('/sys/block/{}/size'.format(dos_path.split('/')[-1])) as dev:
            i_size = int(dev.read())*512
            is_dev=True
    else:
        i_size = os.stat(dos_path).st_size
    # fs_data_size = i_size // 512 - ROOT_START_IDX * 1024 * 1024 // 512
    # fat_size = (fs_data_size) // CLUS_SIZE // 128
    
    print('磁盘占用{}字节'.format(i_size))

    fat_size=(i_size//512-20480)//(128*CLUS_SIZE+2)
    # fat_size = ((i_size-RESERVE_SIZE*512)//(CLUS_SIZE//4+2))//512
    print('FAT表占用{}扇区'.format(fat_size))
    fs_data_size = i_size//512 - RESERVE_SIZE
    print('文件系统占用{}扇区'.format(fs_data_size))
    if is_dev:
        if is_mount(dos_path):
            ret=os.system('umount {}'.format(dos_path))
            if ret!=0:
                exit(-1)


    fd = os.open(dos_path, os.O_RDWR)
    # ===============DBR===============
    pos=os.lseek(fd,0,os.SEEK_CUR)
    print('DBR初始化，位置[{}]'.format(pos))
    data = b'\xeb\x58\x90' #跳转指令
    data = data + b'MSDOS5.0' #文件系统标志和版本号，这里为MSDOC5.0
    data = data + int(512).to_bytes(2, 'little') #每扇区字节数
    data = data + int(CLUS_SIZE).to_bytes(1, 'little') #每簇扇区数
    data = data + int(RESERVE_SIZE).to_bytes(2, 'little') #保留扇区数
    data = data + b'\x02' #FAT表个数
    data = data + b'\x00\x00' #FAT32必须等于0，FAT12/FAT16为根目录中目录的个数
    data = data + b'\x00\x00' #FAT32必须等于0，FAT12/FAT16为扇区总数
    data = data + b'\xf0' #哪种存储介质，0xF8标准值，可移动存储介质，常用的 0xF0
    data = data + b'\x00\x00' #FAT32必须为0，FAT12/FAT16为一个FAT 表所占的扇区数
    data = data + int(0x3f).to_bytes(2, 'little') #每磁道扇区数，只对于有“特殊形状”（由磁头和柱面每 分割为若干磁道）的存储介质有效，63（0x00 3F）
    data = data + int(0xff).to_bytes(2, 'little') #磁头数，只对特殊的介质才有效，255（0x00 FF）
    data = data + int(0).to_bytes(4, 'little') #EBR分区之前所隐藏的扇区数，8192（0x00 00 20 00），与MBR中地址0x1C6开始的4个字节数值相等
    data = data + int(i_size/512).to_bytes(4, 'little') #文件系统总扇区数
    data = data + int(fat_size).to_bytes(4, 'little') #每个FAT表占用扇区数
    data = data + b'\x00\x00' #标记，此域FAT32 特有
    data = data + b'\x00\x00' #FAT32版本号0.0，FAT32特有
    data = data + int(2).to_bytes(4, 'little') #根目录所在第一个簇的簇号，（虽然在FAT32文件系统 下，根目录可以存放在数据区的任何位置，但是通常情况下还是起始于2号簇）
    data = data + int(1).to_bytes(2, 'little') #FSINFO（文件系统信息扇区）扇区号1，该扇区为操作 系统提供关于空簇总数及下一可用簇的信息
    data = data + int(6).to_bytes(2, 'little') #备份引导扇区的位置。备份引导扇区总是位于文件系统 的6号扇区
    data = data + b'\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00' #用于以后FAT 扩展使用
    data = data + b'\x80' #与FAT12/16 的定义相同，只不过两者位于启动扇区不同的位置而已
    data = data + b'\x01' #与FAT12/16 的定义相同，只不过两者位于启动扇区不同的位置而已
    data = data + b'\x29' #扩展引导标志，0x29。与FAT12/16 的定义相同，只不过 两者位于启动扇区不同的位置而已
    data = data + int(random.randint(1, 0xffffffff)).to_bytes(4, 'little') #卷序列号。通常为一个随机值
    data = data + b'NO NAME    ' #卷标（ASCII码），如果建立文件系统的时候指定了卷 标，会保存在此
    data = data + b'FAT32   ' #文件系统格式的ASCII码，FAT32
    data = data + (b'\x00' * 420) #90~509共420字节，未使用。该部分没有明确的用途
    data = data + b'\x55\xaa' #签名标志“55 AA”
    MBR=data
    os.write(fd, data)
    pos=os.lseek(fd,0,os.SEEK_CUR)
    print('DBR初始化成功，位置[{}]'.format(pos)) # 30705087 扇区

    # ===============FSINFO===============
    pos=os.lseek(fd,0,os.SEEK_CUR)
    print('FSINFO初始化,位置[{}]'.format(pos))
    data = b'\x52\x52\x61\x41'
    data = data + (b'\x00' * 480)
    data = data + b'\x72\x72\x41\x61'
    data = data + int(fs_data_size//CLUS_SIZE - 1).to_bytes(4, 'little')
    data = data + int(3).to_bytes(4, 'little')
    data = data + (b'\x00' * 14)
    data = data + b'\x55\xaa'
    os.write(fd, data)
    pos=os.lseek(fd,0,os.SEEK_CUR)
    print('FSINFO初始化成功,位置[{}]'.format(fs_data_size//CLUS_SIZE - 3))

    # ===============Backup MBR===============
    os.lseek(fd,6 * 512,os.SEEK_SET) # 6号扇区
    pos=os.lseek(fd,0,os.SEEK_CUR)
    print('MBR备份,位置[{}]'.format(pos))
    os.write(fd, MBR)
    pos=os.lseek(fd,0,os.SEEK_CUR)
    print('MBR备份成功,位置[{}]'.format(pos))

    # ===============FAT1===============
    os.lseek(fd,RESERVE_SIZE * 512,os.SEEK_SET)
    pos=os.lseek(fd,0,os.SEEK_CUR)
    print('FAT1表初始化,位置[{}]'.format(pos))
    fill_fat(fd, fat_size, tp)
    pos=os.lseek(fd,0,os.SEEK_CUR)
    print('FAT1表初始化成功,位置[{}]'.format(pos))
    # ===============FAT2===============
    os.lseek(fd,(RESERVE_SIZE+fat_size)*512,os.SEEK_SET)
    pos=os.lseek(fd,0,os.SEEK_CUR)
    print('FAT2表初始化，位置[{}]'.format(pos))
    pos=os.lseek(fd,0,os.SEEK_CUR)
    fill_fat(fd, fat_size, tp)
    pos=os.lseek(fd,0,os.SEEK_CUR)
    print('FAT2表初始化成功,位置[{}]'.format(pos))

    # ===============根目录===============
    os.lseek(fd,(RESERVE_SIZE+fat_size*2)*512,os.SEEK_SET)
    pos=os.lseek(fd,0,os.SEEK_CUR)
    print('根目录初始化,位置[{}]'.format(pos))
    fill_root_dir(fd)
    pos=os.lseek(fd,0,os.SEEK_CUR)
    print('根目录初始化成功,位置[{}]'.format(pos))

    print('开始写回磁盘')
    os.fsync(fd)
    print('同步磁盘成功')
    
    os.close(fd)
    print('格式化完成！！！')


if __name__ == '__main__':
    if (len(sys.argv) <= 1):
        print("Error: 请输入需要格式化的路径")
        exit(-1)
    
    
    write_dos(sys.argv[1], SPEED_FILL.MID)
